import torch
import deep_Q_network as dqn
from parallel_upgrades_env import env
import graph_neural_network as GNN


# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")
agent = dqn.GNNDQNAgent(graph_dim = 2 * max(env.max_versions) + 1,
    graph_hid_dim   = 32,
    q_net_hid_dim   = 64,
    agent_feat_dim  = env.feats_dim,        # 环境里 get_state_action_features 返回的 D)
    device          = device)
vdg = GNN.VersionedDependencyGraph(env.dependency_matrix, env.jump_matrix, env.upgrade_time, device)
dqn.load_agent(agent, './gnn_dqn_agent.pth', map_location=device)

done = False
cur_versions = torch.tensor(env.state, dtype=torch.long, device=device)
graph_data = vdg.build_graph(cur_versions)
valid_idxs, feats = env.get_state_action_features()
feats = torch.tensor(feats, device=device)

while not done:
    pick = agent.select_action(graph_data, feats, is_training=False)
    action = valid_idxs[pick]
    _, _, done, _ = env.step(action, is_need_visualization=True)
    env.render(done)

    # update for next step
    cur_versions = torch.tensor(env.state, dtype=torch.long, device=device)
    graph_data = vdg.build_graph(cur_versions)
    valid_idxs, feats = env.get_state_action_features()
    feats = torch.tensor(feats, device=device)

